import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
# TODO: solve no qt platform plugin could be initialized
matplotlib.use('Agg')

feature = np.load('feature_test.npy')
label = np.load('label.npy')
print(feature.shape) #(11396, 768)
print(label.shape) #(11396, 768)
categories = ['ADI_particle_developed', 'Array_peeling', 'Cu_missing', 'Other_peeling', 'Partial_etch',
                  'Pattern_fail',
                  'PR_peeling', 'Seam', 'False_defect', 'Surface_particle', 'Burried_particle', 'Cu_diffuse',
                  'Prelayer_defect_developed',
                  'Void', 'Residue', 'Scratch']

categories_to_id = dict((c, i) for i, c in enumerate(categories))
id_to_categories = dict((v, k) for k, v in categories_to_id.items())

# tsne = TSNE(n_components=2, random_state=42)
tsne = TSNE(n_components=2, n_iter=20000)
tsne_results = tsne.fit_transform(feature)

markers = ['o', 's', '^', 'v', '>', '<', 'p', '*', 'h', 'H', 'D', 'd', 'P', 'X', '8', '|']
cmap = plt.cm.get_cmap('tab20', 16)
colors = cmap.colors
print(colors)

plt.figure(figsize=(16, 16))
for i in range(len(categories)):
    indices = label == i
    plt.scatter(tsne_results[indices, 0], tsne_results[indices, 1], label=str(id_to_categories[i]),
                c=colors[i], marker=markers[i], alpha=0.7, edgecolors='w', linewidths=0.5)

plt.legend()
plt.title('t-sne visualization of Defect test set')
plt.xlabel('t-sne feature 1')
plt.ylabel('t-sne feature 2')
plt.savefig('t-sne.pdf')